#pragma once

#include "koopa.h"
#include <algorithm>
#include <array>
#include <cassert>
#include <iostream>
#include <map>
#include <memory>
#include <sstream>
#include <string>

const std::map<int, std::string> REG2NAME{
    {0, "zero"}, {1, "ra"},  {2, "sp"},   {3, "gp"},   {4, "tp"},  {5, "t0"},
    {6, "t1"},   {7, "t2"},  {8, "s0"},   {9, "s1"},   {10, "a0"}, {11, "a1"},
    {12, "a2"},  {13, "a3"}, {14, "a4"},  {15, "a5"},  {16, "a6"}, {17, "a7"},
    {18, "s2"},  {19, "s3"}, {20, "s4"},  {21, "s5"},  {22, "s6"}, {23, "s7"},
    {24, "s8"},  {25, "s9"}, {26, "s10"}, {27, "s11"}, {28, "t3"}, {29, "t4"},
    {30, "t5"},  {31, "t6"},
};

const std::map<std::string, int> NAME2REG{
    {"zero", 0}, {"ra", 1},  {"sp", 2},   {"gp", 3},   {"tp", 4},  {"t0", 5},
    {"t1", 6},   {"t2", 7},  {"s0", 8},   {"s1", 9},   {"a0", 10}, {"a1", 11},
    {"a2", 12},  {"a3", 13}, {"a4", 14},  {"a5", 15},  {"a6", 16}, {"a7", 17},
    {"s2", 18},  {"s3", 19}, {"s4", 20},  {"s5", 21},  {"s6", 22}, {"s7", 23},
    {"s8", 24},  {"s9", 25}, {"s10", 26}, {"s11", 27}, {"t3", 28}, {"t4", 29},
    {"t5", 30},  {"t6", 31},
};

class Registers {
    std::array<const koopa_raw_value_kind_t*, 32> regs{nullptr};
    constexpr static std::array<int, 7> Tmpregs{5, 6, 7, 28, 29, 30, 31};
    constexpr static std::array<int, 2> Retregs{10, 11};

    int get_tmpreg()
    {
        for (auto index : Tmpregs) {
            if (regs[index] == nullptr) {
                return index;
            }
        }
        std::cerr << "Temp registers used up." << std::endl;
        assert(false);
    }
    int get_retreg()
    {
        for (auto index : Retregs) {
            if (regs[index] == nullptr) {
                return index;
            }
        }
        std::cerr << "Return registers used up." << std::endl;
        assert(false);
    }

  public:
    enum class StorageType { tmp, ret };
    constexpr static int None = -1;

    int store(const koopa_raw_value_kind_t* kindptr, StorageType sttype)
    {
        int res;
        switch (sttype) {
            case StorageType::tmp: {
                res       = get_tmpreg();
                regs[res] = kindptr;
            } break;
            case StorageType::ret: {
                res       = get_retreg();
                regs[res] = kindptr;
            } break;
            default: assert(false); break;
        }

        return res;
    }

    int store(const koopa_raw_value_kind_t* kindptr, int index)
    {
        regs[index] = kindptr;
        return index;
    }

    void release(int index)
    {
        regs[index] = nullptr;
    }

    const std::string& name(int index)
    {
        return REG2NAME.at(index);
    }

    int find(const koopa_raw_value_kind_t* kindptr)
    {
        for (size_t i = 0; i != regs.size(); i++) {
            // 跳过特殊寄存器, 例如返回值
            if (std::find(Retregs.begin(), Retregs.end(), i) != Retregs.end()) {
                continue;
            }
            if (regs[i] == kindptr) {
                return i;
            }
        }

        return None;
    }
};

extern Registers rv32regs;

std::unique_ptr<std::string>
     to_riscv32_asm(std::unique_ptr<std::string>& koopa_ir_text);
void visit_koopa_bin(const char* str, std::ostringstream&, std::ostringstream&);

void visit_koopa(const koopa_raw_program_t&,
                 std::ostringstream&,
                 std::ostringstream&);
void visit_koopa(const koopa_raw_slice_t&,
                 std::ostringstream&,
                 std::ostringstream&);
void visit_koopa(const koopa_raw_function_t&,
                 std::ostringstream&,
                 std::ostringstream&);
void visit_koopa(const koopa_raw_basic_block_t&,
                 std::ostringstream&,
                 std::ostringstream&);
void visit_koopa(const koopa_raw_value_t&,
                 std::ostringstream&,
                 std::ostringstream&);
void visit_koopa(const koopa_raw_return_t&,
                 std::ostringstream&,
                 std::ostringstream&);
void visit_koopa(const koopa_raw_binary_t&,
                 std::ostringstream&,
                 std::ostringstream&,
                 const koopa_raw_value_kind_t* kindptr);